#include "pythonObjects.h"
#include "utils.h"
#include "pythonXFunctions.h"
#include "exceptions.h"
#include <s2e/S2EExecutionState.h>
#include <s2e/S2E.h>
#include <s2e/cpu.h>

namespace s2e {
  namespace plugins {
    namespace python {

void S2EState_dealloc(PyObject *self) {
	std::cout << std::flush;
	Py_TYPE(self)->tp_free((PyObject*) self);
}
static int S2EState_print(PyObject *self, void *io, int s) {
	std::cout << "id=" << PyS2EState_AsS2EExecutionState(self)->getID();
	return 0;
}

static PyObject* S2EState_str(PyObject *self) {
	try {
		std::stringstream str;
		str << PyS2EState_AsS2EExecutionState(self)->getID();
		return PyStr_FromFormat("id=%s", str.str().c_str());
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}
static PyObject* S2EState_getPointerSize(PyObject *self, PyObject *noarg) {
	try {
		s2e::uint32 size =
				PyS2EState_AsS2EExecutionState(self)->getPointerSize();
		return PyLong_FromUint32(size);
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}

static PyObject* S2EState_kill(PyObject *self, PyObject *message) {
	try {
		if (!PyStr_Check(message))
			return PyErr_Format(PyExc_TypeError,
					"S2EState::kill(msg): Expected a string as argument.");
		std::stringstream ss;
		ss << "PyS2eState: killed message:" << message;
		g_s2e->getExecutor()->terminateState(
				*PyS2EState_AsS2EExecutionState(self), ss.str());
		Py_INCREF(Py_None);
		return Py_None;
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}

static PyObject* S2EState_createSymbolicValue(PyObject *self, PyObject *args) {
	try {
		PyObject *concrate_data = nullptr;
		PyObject *name = nullptr;
		std::string cname = "";

		/* Extract arguments */
		if (PyArg_ParseTuple(args, "|OO", &concrate_data, &name) == false) {
			return PyErr_Format(PyExc_TypeError,
					"S2EState::createSymbolicValue(): Invalid number of arguments");
		}
		if (concrate_data != nullptr || !PyBytes_Check(concrate_data))
			return PyErr_Format(PyExc_TypeError,
					"S2EState::createSymbolicValue(): Expected bytes as first argument.");

		if (name == nullptr && (!PyStr_Check(name)))
			return PyErr_Format(PyExc_TypeError,
					"S2EState::createSymbolicValue(): Expected a string as second argument.");

		if (name != nullptr)
			cname = PyStr_AsString(name);

		s2e::uint8 *opcode = reinterpret_cast<s2e::uint8*>(PyBytes_AsString(
				concrate_data));
		s2e::uint32 size = static_cast<s2e::uint32>(PyBytes_Size(concrate_data));
		std::vector<uint8_t> buffer(size);
		for (unsigned i = 0; i < size; ++i) {
			buffer[i] = opcode[i];
		}
		S2EExecutionState *state = PyS2EState_AsS2EExecutionState(self);
		klee::ref<klee::Expr> value = state->createSymbolicValue(cname,
				size * 8, buffer);
		g_s2e->getDebugStream(state) << "S2EState createSymbolicValue: "
				<< value << "\n";

		Py_INCREF(Py_None);
		return Py_None;
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}



static uint64_t fetchRegister(S2EExecutionState *state, uint8_t regIdx,
		bool isNormal) {
	target_ulong reg;
	if (isNormal) {
		reg = state->regs()->read<target_ulong>(CPU_OFFSET(regs[regIdx]));
	} else { // segment
		reg = state->regs()->read<target_ulong>(CPU_OFFSET(segs[regIdx]));
	}
	return reg;
}
      /*
       * #define CPU_NB_REGS64 16
	   * #define CPU_NB_REGS32 8
	   * #define R_EAX 0
	   * #define R_ECX 1
	   * #define R_EDX 2
	   * #define R_EBX 3
	   * #define R_ESP 4
	   * #define R_EBP 5
	   * #define R_ESI 6
	   * #define R_EDI 7
       */
static PyObject* S2EState_getRegisters(PyObject *self, PyObject *noarg) {
	try {
		const auto state = PyS2EState_AsS2EExecutionState(self);
		PyObject *dict = xPyDict_New();
		xPyDict_SetItem(dict, PyStr_FromString("rax"),
				PyLong_FromUint64(fetchRegister(state, 0, true)));
		xPyDict_SetItem(dict, PyStr_FromString("rbx"),
				PyLong_FromUint64(fetchRegister(state, 3, true)));
		xPyDict_SetItem(dict, PyStr_FromString("rcx"),
				PyLong_FromUint64(fetchRegister(state, 1, true)));
		xPyDict_SetItem(dict, PyStr_FromString("rdx"),
				PyLong_FromUint64(fetchRegister(state, 2, true)));
		xPyDict_SetItem(dict, PyStr_FromString("rdi"),
				PyLong_FromUint64(fetchRegister(state, 7, true)));
		xPyDict_SetItem(dict, PyStr_FromString("rsi"),
				PyLong_FromUint64(fetchRegister(state, 6, true)));
		xPyDict_SetItem(dict, PyStr_FromString("rbp"),
				PyLong_FromUint64(fetchRegister(state, 5, true)));
		xPyDict_SetItem(dict, PyStr_FromString("rsp"),
				PyLong_FromUint64(fetchRegister(state, 4, true)));
		xPyDict_SetItem(dict, PyStr_FromString("r8"),
				PyLong_FromUint64(fetchRegister(state, 8, true)));
		xPyDict_SetItem(dict, PyStr_FromString("r9"),
				PyLong_FromUint64(fetchRegister(state, 9, true)));
		xPyDict_SetItem(dict, PyStr_FromString("r10"),
				PyLong_FromUint64(fetchRegister(state, 10, true)));
		xPyDict_SetItem(dict, PyStr_FromString("r11"),
				PyLong_FromUint64(fetchRegister(state, 11, true)));
		xPyDict_SetItem(dict, PyStr_FromString("r12"),
				PyLong_FromUint64(fetchRegister(state, 12, true)));
		xPyDict_SetItem(dict, PyStr_FromString("r13"),
				PyLong_FromUint64(fetchRegister(state, 13, true)));
		xPyDict_SetItem(dict, PyStr_FromString("r14"),
				PyLong_FromUint64(fetchRegister(state, 14, true)));
		xPyDict_SetItem(dict, PyStr_FromString("r15"),
				PyLong_FromUint64(fetchRegister(state, 15, true)));

		xPyDict_SetItem(dict, PyStr_FromString("rip"),
				PyLong_FromUint64(state->regs()->getPc()));
		xPyDict_SetItem(dict, PyStr_FromString("eflags"),
				PyLong_FromUint64(state->regs()->getFlags()));
		return dict;
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}
static PyObject* S2EState_getRax(PyObject *self, PyObject *noarg) {
	try {
		uint64_t rax = fetchRegister(PyS2EState_AsS2EExecutionState(self), 0, true);
		return PyLong_FromUint64(rax);
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}
static PyObject* S2EState_getPc(PyObject *self, PyObject *noarg) {
	try {
		uint64_t pc =
		PyS2EState_AsS2EExecutionState(self)->regs()->getPc();
		return PyLong_FromUint64(pc);
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}
static PyObject* S2EState_getSp(PyObject *self, PyObject *noarg) {
	try {
		uint64_t sp =
		PyS2EState_AsS2EExecutionState(self)->regs()->getSp();
		return PyLong_FromUint64(sp);
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}

static PyObject* S2EState_readBytes(PyObject *self, PyObject *args) {
	PyObject *addr = nullptr;
	PyObject *size = nullptr;

	/* Extract arguments */
	if (PyArg_ParseTuple(args, "|OO", &addr, &size) == false) {
		return PyErr_Format(PyExc_TypeError,
				"S2EState::readBytes(): Invalid number of arguments");
	}
	try {
		s2e::uint64 addr_num = PyLong_AsUint64(addr);
		s2e::uint32 size_num = static_cast<s2e::uint32>(PyLong_AsUsize(size));

		std::vector<uint8_t> buffer(size_num);

		if (!PyS2EState_AsS2EExecutionState(self)->mem()->read(addr_num,
				buffer.data(), size_num * sizeof(uint8_t))) {
			Py_INCREF(Py_None);
			return Py_None;
		}
		return PyBytes_FromStringAndSize(
				reinterpret_cast<const char*>(buffer.data()), size_num);
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}
static PyObject* S2EState_readString(PyObject *self, PyObject *args) {
	PyObject *addr = nullptr;
	PyObject *size = nullptr;

	/* Extract arguments */
	if (PyArg_ParseTuple(args, "|OO", &addr, &size) == false) {
		return PyErr_Format(PyExc_TypeError,
				"S2EState::readBytes(): Invalid number of arguments");
	}
	try {
		s2e::uint64 addr_num = PyLong_AsUint64(addr);
		s2e::uint32 size_num = static_cast<s2e::uint32>(PyLong_AsUsize(size));
    	std::string target_str;
		if (!PyS2EState_AsS2EExecutionState(self)->mem()->readString(addr_num,
				target_str, size_num * sizeof(uint8_t))) {
			Py_INCREF(Py_None);
			return Py_None;
		}
		return PyBytes_FromStringAndSize(
						reinterpret_cast<const char*>(target_str.c_str()), target_str.size());
		//return PyStr_FromString(target_str.c_str());
	} catch (const s2e::exceptions::Exception &e) {
		return PyErr_Format(PyExc_TypeError, "%s", e.what());
	}
}
      //! S2EState methods.
      PyMethodDef S2EState_callbacks[] = {
    		  {"getPointerSize",            S2EState_getPointerSize,           	METH_NOARGS,      ""},
			  {"kill",              		S2EState_kill,                	  	METH_O,  	     ""},
			  {"createSymbolicValue",       S2EState_createSymbolicValue,     	METH_VARARGS,     ""},
    		  {"getRegisters",              S2EState_getRegisters,            	METH_NOARGS,      ""},
			  {"getRax",           			S2EState_getRax,           			METH_NOARGS,      ""},
			  {"getSp",           			S2EState_getSp,           			METH_NOARGS,      ""},
			  {"getPc",           			S2EState_getPc,          			METH_NOARGS,      ""},
			  {"readBytes",       			S2EState_readBytes,     			METH_VARARGS,     ""},
			  {"readString",       			S2EState_readString,     			METH_VARARGS,     ""},
			  {nullptr,                     nullptr,                              0,               nullptr}
      };


      PyTypeObject S2EState_Type = {
        PyVarObject_HEAD_INIT(&PyType_Type, 0)
        "S2EState",                              /* tp_name */
        sizeof(S2EState_Object),                 /* tp_basicsize */
        0,                                          /* tp_itemsize */
        (destructor)S2EState_dealloc,            /* tp_dealloc */
        (printfunc)S2EState_print,               /* tp_print */
        0,                                          /* tp_getattr */
        0,                                          /* tp_setattr */
        0,                                          /* tp_compare */
        (reprfunc)S2EState_str,                  /* tp_repr */
        0,                                          /* tp_as_number */
        0,                                          /* tp_as_sequence */
        0,                                          /* tp_as_mapping */
        0,                                          /* tp_hash */
        0,                                          /* tp_call */
        (reprfunc)S2EState_str,                  /* tp_str */
        0,                                          /* tp_getattro */
        0,                                          /* tp_setattro */
        0,                                          /* tp_as_buffer */
        Py_TPFLAGS_DEFAULT,                         /* tp_flags */
        "S2EState objects",                      /* tp_doc */
        0,                                          /* tp_traverse */
        0,                                          /* tp_clear */
        0,                                          /* tp_richcompare */
        0,                                          /* tp_weaklistoffset */
        0,                                          /* tp_iter */
        0,                                          /* tp_iternext */
        S2EState_callbacks,                      /* tp_methods */
        0,                                          /* tp_members */
        0,                                          /* tp_getset */
        0,                                          /* tp_base */
        0,                                          /* tp_dict */
        0,                                          /* tp_descr_get */
        0,                                          /* tp_descr_set */
        0,                                          /* tp_dictoffset */
        0,                                          /* tp_init */
        0,                                          /* tp_alloc */
        0,                                          /* tp_new */
        0,                                          /* tp_free */
        0,                                          /* tp_is_gc */
        0,                                          /* tp_bases */
        0,                                          /* tp_mro */
        0,                                          /* tp_cache */
        0,                                          /* tp_subclasses */
        0,                                          /* tp_weaklist */
        0,                                          /* tp_del */
        0,                                          /* tp_version_tag */
        0,                                          /* tp_finalize */
      };


      PyObject* PyS2EState(s2e::S2EExecutionState* state) {
    	  S2EState_Object* object;

        PyType_Ready(&S2EState_Type);
        object = PyObject_NEW(S2EState_Object, &S2EState_Type);
        if (object != NULL)
          object->state = state;

        return (PyObject*)object;
      }

    };
  };
};
